# 仿真环境初始化模块
import torch
import ai2thor.controller

# torch.cuda.set_per_process_memory_fraction(0.8)  # 单卡使用不超过80%显存

class ThorEnvBuilder:
    def __init__(self, config_path):
        self.load_config(config_path)
        self.controller = None

    def load_config(self, path):
        """加载YAML配置文件"""
        with open(path) as f:
            self.config = yaml.safe_load(f)

    def initialize(self):
        """创建AI2-THOR控制器实例"""
        self.controller = ai2thor.controller.Controller(
            width=self.config['render_settings']['width'],
            height=self.config['render_settings']['height'],
            renderDepthImage=self.config['render_settings']['render_depth'],
            gridSize=self.config['agent_params']['grid_size']
        )

    def reset_scene(self, scene_name):
        """重置到指定场景"""
        self.controller.reset(scene_name)
        self._apply_texture_variation()

    def _apply_texture_variation(self):
        """随机应用材质变化"""
        texture = random.choice(self.config['scene_variations']['texture_packs'])
        for obj in self.controller.last_event.metadata["objects"]:
            if obj['objectType'] in ['Wall', 'Floor']:
                self.controller.step(
                    action='ChangeMaterial',
                    objectId=obj['objectId'],
                    texturePack=texture
                )